from flask import Flask, request, jsonify, send_file
import os
import json
from datetime import datetime
import uuid

app = Flask(__name__)

# 邮件存储目录
MAIL_STORAGE = 'mails'
ATTACHMENT_STORAGE = 'attachments'

# 确保存储目录存在
os.makedirs(MAIL_STORAGE, exist_ok=True)
os.makedirs(ATTACHMENT_STORAGE, exist_ok=True)

# 两个用户的标识
USER1 = 'user1'
USER2 = 'user2'

@app.route('/send', methods=['POST'])
def send_mail():
    # 获取请求数据
    data = request.form.to_dict()
    sender = data.get('sender')
    content = data.get('content')
    subject = data.get('subject', '无主题')

    # 验证发送者
    if sender not in [USER1, USER2]:
        return jsonify({'status': 'error', 'message': '无效的发送者'}), 400

    # 确定接收者
    receiver = USER2 if sender == USER1 else USER1

    # 生成邮件ID
    mail_id = str(uuid.uuid4())

    # 处理附件
    attachments = []
    if 'attachment' in request.files:
        file = request.files['attachment']
        if file.filename:
            # 保存附件
            attachment_id = str(uuid.uuid4())
            file_path = os.path.join(ATTACHMENT_STORAGE, attachment_id)
            file.save(file_path)
            attachments.append({
                'id': attachment_id,
                'name': file.filename,
                'path': file_path
            })

    # 创建邮件对象
    mail = {
        'id': mail_id,
        'sender': sender,
        'receiver': receiver,
        'subject': subject,
        'content': content,
        'timestamp': datetime.now().isoformat(),
        'read': False,
        'attachments': attachments
    }

    # 保存邮件
    mail_file = os.path.join(MAIL_STORAGE, f'{mail_id}.json')
    with open(mail_file, 'w', encoding='utf-8') as f:
        json.dump(mail, f, ensure_ascii=False, indent=2)

    return jsonify({'status': 'success', 'message': '邮件发送成功', 'mail_id': mail_id})

@app.route('/inbox/<user>', methods=['GET'])
def get_inbox(user):
    # 验证用户
    if user not in [USER1, USER2]:
        return jsonify({'status': 'error', 'message': '无效的用户'}), 400

    # 获取参数
    unread_only = request.args.get('unread_only', 'false').lower() == 'true'

    # 获取收件箱邮件
    mails = []
    for filename in os.listdir(MAIL_STORAGE):
        if filename.endswith('.json'):
            with open(os.path.join(MAIL_STORAGE, filename), 'r', encoding='utf-8') as f:
                mail = json.load(f)
                if mail['receiver'] == user:
                    if not unread_only or not mail['read']:
                        mails.append(mail)

    # 按时间排序
    mails.sort(key=lambda x: x['timestamp'], reverse=True)

    return jsonify({'status': 'success', 'mails': mails})

@app.route('/mark_as_read/<mail_id>', methods=['POST'])
def mark_as_read(mail_id):
    # 查找邮件
    mail_file = os.path.join(MAIL_STORAGE, f'{mail_id}.json')
    if not os.path.exists(mail_file):
        return jsonify({'status': 'error', 'message': '邮件不存在'}), 404

    # 读取邮件
    with open(mail_file, 'r', encoding='utf-8') as f:
        mail = json.load(f)

    # 更新状态
    mail['read'] = True

    # 保存邮件
    with open(mail_file, 'w', encoding='utf-8') as f:
        json.dump(mail, f, ensure_ascii=False, indent=2)

    return jsonify({'status': 'success', 'message': '邮件已标记为已读'})

@app.route('/attachment/<attachment_id>', methods=['GET'])
def get_attachment(attachment_id):
    # 查找附件
    attachment_path = os.path.join(ATTACHMENT_STORAGE, attachment_id)
    if not os.path.exists(attachment_path):
        return jsonify({'status': 'error', 'message': '附件不存在'}), 404

    # 查找对应的邮件，获取附件名称
    attachment_name = 'attachment'
    for filename in os.listdir(MAIL_STORAGE):
        if filename.endswith('.json'):
            with open(os.path.join(MAIL_STORAGE, filename), 'r', encoding='utf-8') as f:
                mail = json.load(f)
                for att in mail.get('attachments', []):
                    if att['id'] == attachment_id:
                        attachment_name = att['name']
                        break

    # 返回附件
    return send_file(attachment_path, as_attachment=True, download_name=attachment_name)

if __name__ == '__main__':
    app.run(debug=False, host='0.0.0.0', port=8000)
